// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>

#include <nebula/compute/kernels/codegen_internal.h>
#include <nebula/core/visit_data_inline.h>

namespace nebula::compute::internal {

    struct KeyEncoder {
        // the first byte of an encoded key is used to indicate nullity
        static constexpr bool kExtraByteForNull = true;

        static constexpr uint8_t kNullByte = 1;
        static constexpr uint8_t kValidByte = 0;

        virtual ~KeyEncoder() = default;

        virtual void AddLength(const ExecValue &value, int64_t batch_length,
                               int32_t *lengths) = 0;

        virtual void AddLengthNull(int32_t *length) = 0;

        virtual turbo::Status Encode(const ExecValue &, int64_t batch_length,
                                     uint8_t **encoded_bytes) = 0;

        virtual void EncodeNull(uint8_t **encoded_bytes) = 0;

        virtual turbo::Result<std::shared_ptr<ArrayData>> Decode(uint8_t **encoded_bytes,
                                                                 int32_t length, MemoryPool *) = 0;

        // extract the null bitmap from the leading nullity bytes of encoded keys
        static turbo::Status DecodeNulls(MemoryPool *pool, int32_t length, uint8_t **encoded_bytes,
                                         std::shared_ptr<Buffer> *null_bitmap, int32_t *null_count);

        static bool is_null(const uint8_t *encoded_bytes) {
            return encoded_bytes[0] == kNullByte;
        }
    };

    struct BooleanKeyEncoder : KeyEncoder {
        static constexpr int kByteWidth = 1;

        void AddLength(const ExecValue &data, int64_t batch_length, int32_t *lengths) override;

        void AddLengthNull(int32_t *length) override;

        turbo::Status Encode(const ExecValue &data, int64_t batch_length,
                             uint8_t **encoded_bytes) override;

        void EncodeNull(uint8_t **encoded_bytes) override;

        turbo::Result<std::shared_ptr<ArrayData>> Decode(uint8_t **encoded_bytes, int32_t length,
                                                         MemoryPool *pool) override;
    };

    struct FixedWidthKeyEncoder : KeyEncoder {
        explicit FixedWidthKeyEncoder(std::shared_ptr<DataType> type)
                : type_(std::move(type)),
                  byte_width_(turbo::checked_cast<const FixedWidthType &>(*type_).bit_width() / 8) {}

        void AddLength(const ExecValue &data, int64_t batch_length, int32_t *lengths) override;

        void AddLengthNull(int32_t *length) override;

        turbo::Status Encode(const ExecValue &data, int64_t batch_length,
                             uint8_t **encoded_bytes) override;

        void EncodeNull(uint8_t **encoded_bytes) override;

        turbo::Result<std::shared_ptr<ArrayData>> Decode(uint8_t **encoded_bytes, int32_t length,
                                                         MemoryPool *pool) override;

        std::shared_ptr<DataType> type_;
        int byte_width_;
    };

    struct DictionaryKeyEncoder : FixedWidthKeyEncoder {
        DictionaryKeyEncoder(std::shared_ptr<DataType> type, MemoryPool *pool)
                : FixedWidthKeyEncoder(std::move(type)), pool_(pool) {}

        turbo::Status Encode(const ExecValue &data, int64_t batch_length,
                             uint8_t **encoded_bytes) override;

        turbo::Result<std::shared_ptr<ArrayData>> Decode(uint8_t **encoded_bytes, int32_t length,
                                                         MemoryPool *pool) override;

        MemoryPool *pool_;
        std::shared_ptr<Array> dictionary_;
    };

    template<typename T>
    struct VarLengthKeyEncoder : KeyEncoder {
        using Offset = typename T::offset_type;

        void AddLength(const ExecValue &data, int64_t batch_length, int32_t *lengths) override {
            if (data.is_array()) {
                int64_t i = 0;
                VisitArraySpanInline<T>(
                        data.array,
                        [&](std::string_view bytes) {
                            lengths[i++] +=
                                    kExtraByteForNull + sizeof(Offset) + static_cast<int32_t>(bytes.size());
                        },
                        [&] { lengths[i++] += kExtraByteForNull + sizeof(Offset); });
            } else {
                const Scalar &scalar = *data.scalar;
                const int32_t buffer_size =
                        scalar.is_valid ? static_cast<int32_t>(UnboxScalar<T>::Unbox(scalar).size())
                                        : 0;
                for (int64_t i = 0; i < batch_length; i++) {
                    lengths[i] += kExtraByteForNull + sizeof(Offset) + buffer_size;
                }
            }
        }

        void AddLengthNull(int32_t *length) override {
            *length += kExtraByteForNull + sizeof(Offset);
        }

        turbo::Status Encode(const ExecValue &data, int64_t batch_length,
                             uint8_t **encoded_bytes) override {
            if (data.is_array()) {
                VisitArraySpanInline<T>(
                        data.array,
                        [&](std::string_view bytes) {
                            auto &encoded_ptr = *encoded_bytes++;
                            *encoded_ptr++ = kValidByte;
                            turbo::safe_store(encoded_ptr, static_cast<Offset>(bytes.size()));
                            encoded_ptr += sizeof(Offset);
                            memcpy(encoded_ptr, bytes.data(), bytes.size());
                            encoded_ptr += bytes.size();
                        },
                        [&] {
                            auto &encoded_ptr = *encoded_bytes++;
                            *encoded_ptr++ = kNullByte;
                            turbo::safe_store(encoded_ptr, static_cast<Offset>(0));
                            encoded_ptr += sizeof(Offset);
                        });
            } else {
                const auto &scalar = data.scalar_as<BaseBinaryScalar>();
                if (scalar.is_valid) {
                    const auto &bytes = *scalar.value;
                    for (int64_t i = 0; i < batch_length; i++) {
                        auto &encoded_ptr = *encoded_bytes++;
                        *encoded_ptr++ = kValidByte;
                        turbo::safe_store(encoded_ptr, static_cast<Offset>(bytes.size()));
                        encoded_ptr += sizeof(Offset);
                        memcpy(encoded_ptr, bytes.data(), bytes.size());
                        encoded_ptr += bytes.size();
                    }
                } else {
                    for (int64_t i = 0; i < batch_length; i++) {
                        auto &encoded_ptr = *encoded_bytes++;
                        *encoded_ptr++ = kNullByte;
                        turbo::safe_store(encoded_ptr, static_cast<Offset>(0));
                        encoded_ptr += sizeof(Offset);
                    }
                }
            }
            return turbo::OkStatus();
        }

        void EncodeNull(uint8_t **encoded_bytes) override {
            auto &encoded_ptr = *encoded_bytes;
            *encoded_ptr++ = kNullByte;
            turbo::safe_store(encoded_ptr, static_cast<Offset>(0));
            encoded_ptr += sizeof(Offset);
        }

        turbo::Result<std::shared_ptr<ArrayData>> Decode(uint8_t **encoded_bytes, int32_t length,
                                                         MemoryPool *pool) override {
            std::shared_ptr<Buffer> null_buf;
            int32_t null_count;
            TURBO_RETURN_NOT_OK(DecodeNulls(pool, length, encoded_bytes, &null_buf, &null_count));

            Offset length_sum = 0;
            for (int32_t i = 0; i < length; ++i) {
                length_sum += turbo::safe_load_as<Offset>(encoded_bytes[i]);
            }

            TURBO_MOVE_OR_RAISE(auto offset_buf,
                                allocate_buffer(sizeof(Offset) * (1 + length), pool));
            TURBO_MOVE_OR_RAISE(auto key_buf, allocate_buffer(length_sum));

            auto raw_offsets = reinterpret_cast<Offset *>(offset_buf->mutable_data());
            auto raw_keys = key_buf->mutable_data();

            Offset current_offset = 0;
            for (int32_t i = 0; i < length; ++i) {
                raw_offsets[i] = current_offset;

                auto key_length = turbo::safe_load_as<Offset>(encoded_bytes[i]);
                encoded_bytes[i] += sizeof(Offset);

                memcpy(raw_keys + current_offset, encoded_bytes[i], key_length);
                encoded_bytes[i] += key_length;

                current_offset += key_length;
            }
            raw_offsets[length] = current_offset;

            return ArrayData::create(
                    type_, length, {std::move(null_buf), std::move(offset_buf), std::move(key_buf)},
                    null_count);
        }

        explicit VarLengthKeyEncoder(std::shared_ptr<DataType> type) : type_(std::move(type)) {}

        std::shared_ptr<DataType> type_;
    };

    struct NullKeyEncoder : KeyEncoder {
        void AddLength(const ExecValue &, int64_t batch_length, int32_t *lengths) override {}

        void AddLengthNull(int32_t *length) override {}

        turbo::Status Encode(const ExecValue &data, int64_t batch_length,
                             uint8_t **encoded_bytes) override {
            return turbo::OkStatus();
        }

        void EncodeNull(uint8_t **encoded_bytes) override {}

        turbo::Result<std::shared_ptr<ArrayData>> Decode(uint8_t **encoded_bytes, int32_t length,
                                                         MemoryPool *pool) override {
            return ArrayData::create(null(), length, {nullptr}, length);
        }
    };

    class TURBO_EXPORT RowEncoder {
    public:
        static constexpr int kRowIdForNulls() { return -1; }

        void init(const std::vector<TypeHolder> &column_types, ExecContext *ctx);

        void Clear();

        turbo::Status EncodeAndAppend(const ExecSpan &batch);

        turbo::Result<ExecBatch> Decode(int64_t num_rows, const int32_t *row_ids);

        inline std::string encoded_row(int32_t i) const {
            if (i == kRowIdForNulls()) {
                return std::string(reinterpret_cast<const char *>(encoded_nulls_.data()),
                                   encoded_nulls_.size());
            }
            int32_t row_length = offsets_[i + 1] - offsets_[i];
            return std::string(reinterpret_cast<const char *>(bytes_.data() + offsets_[i]),
                               row_length);
        }

        int32_t num_rows() const {
            return offsets_.size() == 0 ? 0 : static_cast<int32_t>(offsets_.size() - 1);
        }

    private:
        ExecContext *ctx_;
        std::vector<std::shared_ptr<KeyEncoder>> encoders_;
        std::vector<int32_t> offsets_;
        std::vector<uint8_t> bytes_;
        std::vector<uint8_t> encoded_nulls_;
        std::vector<std::shared_ptr<ExtensionType>> extension_types_;
    };

}  // namespace nebula::compute::internal
